# supertonic_adapter.py
# Description: Supertonic ONNX TTS adapter implementation
#
import asyncio
from pathlib import Path
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
#
# Third-party Imports
import numpy as np
from loguru import logger
#
# Local Imports
from .base import (
    AudioFormat,
    ProviderStatus,
    TTSCapabilities,
    TTSAdapter,
    TTSRequest,
    TTSResponse,
    VoiceInfo,
)
from ..streaming_audio_writer import AudioNormalizer
from ..tts_exceptions import (
    TTSGenerationError,
    TTSModelLoadError,
    TTSModelNotFoundError,
    TTSProviderInitializationError,
    TTSProviderNotConfiguredError,
    TTSUnsupportedFormatError,
    TTSValidationError,
    TTSVoiceNotFoundError,
)
from ..tts_validation import validate_tts_request
#
#######################################################################################################################
#
# Supertonic Adapter Implementation


class SupertonicOnnxAdapter(TTSAdapter):
    """Adapter for the Supertonic ONNX TTS engine."""

    PROVIDER_KEY = "supertonic"
    SUPPORTED_FORMATS = {AudioFormat.MP3, AudioFormat.WAV}
    SUPPORTED_LANGUAGES = {"en"}
    MAX_TEXT_LENGTH = 15000
    DEFAULT_SAMPLE_RATE = 24000

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__(config)
        cfg = config or {}
        extras = cfg.get("extra_params", {}) or {}

        self.onnx_dir = Path(cfg.get("model_path", "models/supertonic/onnx")).expanduser()
        self.voice_styles_dir = Path(extras.get("voice_styles_dir", "models/supertonic/voice_styles")).expanduser()
        self.default_voice = extras.get("default_voice", "supertonic_m1")
        self.voice_files: Dict[str, str] = extras.get("voice_files", {}) or {}
        self.default_total_step = int(extras.get("default_total_step", 5))
        self.default_speed = float(extras.get("default_speed", 1.05))
        self.n_test = int(extras.get("n_test", 1))
        self.device = cfg.get("device", "cpu")
        self.use_gpu = str(self.device).lower() == "cuda"

        self.sample_rate = int(cfg.get("sample_rate", self.DEFAULT_SAMPLE_RATE))
        # Prefer provider-specific chunk size, fall back to global performance chunk size or default
        stream_chunk = extras.get("stream_chunk_size") or cfg.get("stream_chunk_size")
        try:
            if stream_chunk is None:
                perf_cfg = cfg.get("performance", {}) or {}
                stream_chunk = perf_cfg.get("stream_chunk_size")
            self.stream_chunk_size = int(stream_chunk) if stream_chunk else 8192
        except Exception:
            self.stream_chunk_size = 8192

        self._engine: Optional[Any] = None
        self._load_voice_style = None
        self._voice_to_path: Dict[str, Path] = {}
        self._voice_infos: List[VoiceInfo] = []
        self._audio_normalizer = AudioNormalizer()
        self._engine_lock = asyncio.Lock()

    async def ensure_initialized(self) -> bool:
        """
        Ensure the provider is initialized, propagating TTSErrors so callers see
        clear misconfiguration messages (e.g., missing model or voice files).
        """
        if self._initialized:
            return True

        async with self._init_lock:
            if self._initialized:
                return True

            self._status = ProviderStatus.INITIALIZING
            success = await self.initialize()
            if success:
                self._capabilities = await self.get_capabilities()
                self._status = ProviderStatus.AVAILABLE
                self._initialized = True
            else:
                self._status = ProviderStatus.ERROR
            return success

    # ---------------------------------------------------------------------------------
    # Initialization
    # ---------------------------------------------------------------------------------
    async def initialize(self) -> bool:
        """Load the Supertonic ONNX engine and discover voice styles."""
        if self._initialized:
            return True

        if not self.onnx_dir.exists():
            raise TTSModelNotFoundError(
                f"Supertonic ONNX directory not found at {self.onnx_dir}",
                provider=self.PROVIDER_KEY,
                details={"onnx_dir": str(self.onnx_dir)},
            )

        if not self.voice_styles_dir.exists():
            raise TTSModelNotFoundError(
                f"Supertonic voice styles directory not found at {self.voice_styles_dir}",
                provider=self.PROVIDER_KEY,
                details={"voice_styles_dir": str(self.voice_styles_dir)},
            )

        try:
            from tldw_Server_API.app.core.TTS.vendors import supertonic as vendor
        except ImportError as exc:
            raise TTSModelLoadError(
                "Supertonic vendor package not available",
                provider=self.PROVIDER_KEY,
                details={"error": str(exc)},
            ) from exc

        try:
            engine = await asyncio.to_thread(vendor.load_text_to_speech, str(self.onnx_dir), self.use_gpu)
        except FileNotFoundError as exc:
            raise TTSModelNotFoundError(
                f"Supertonic model assets missing under {self.onnx_dir}",
                provider=self.PROVIDER_KEY,
                details={"onnx_dir": str(self.onnx_dir)},
            ) from exc
        except Exception as exc:
            raise TTSProviderInitializationError(
                "Failed to initialize Supertonic engine",
                provider=self.PROVIDER_KEY,
                details={"error": str(exc)},
            ) from exc

        self._engine = engine
        self._load_voice_style = vendor.load_voice_style

        # Sample rate from engine if available
        try:
            self.sample_rate = int(getattr(self._engine, "sample_rate", self.sample_rate))
        except Exception:
            logger.debug("Supertonic: unable to read sample_rate from engine, using default %s", self.sample_rate)

        self._voice_to_path, self._voice_infos = self._load_voice_mappings()

        # Cache capabilities for quick lookups
        self._capabilities = await self.get_capabilities()
        self._initialized = True
        self._status = ProviderStatus.AVAILABLE
        logger.info(
            "Supertonic adapter initialized (onnx_dir=%s, voices=%d, sample_rate=%s)",
            self.onnx_dir,
            len(self._voice_infos),
            self.sample_rate,
        )
        return True

    def _load_voice_mappings(self) -> Tuple[Dict[str, Path], List[VoiceInfo]]:
        """Build mapping of voice IDs to style files and VoiceInfo entries."""
        voice_map: Dict[str, Path] = {}
        voices: List[VoiceInfo] = []
        missing_default = False

        for voice_id, filename in self.voice_files.items():
            path = (self.voice_styles_dir / filename).expanduser()
            if not path.exists():
                if voice_id == self.default_voice:
                    missing_default = True
                logger.warning("Supertonic voice style missing: voice_id=%s path=%s", voice_id, path)
                continue
            voice_map[voice_id] = path
            voices.append(
                VoiceInfo(
                    id=voice_id,
                    name=voice_id.replace("_", " ").title(),
                    gender="male" if "m" in voice_id.lower() else "female" if "f" in voice_id.lower() else None,
                    language="en",
                    description="Supertonic voice style",
                    styles=["neutral"],
                    use_case=["general"],
                )
            )

        if missing_default or self.default_voice not in voice_map:
            raise TTSModelNotFoundError(
                f"Default Supertonic voice style '{self.default_voice}' not found",
                provider=self.PROVIDER_KEY,
                details={"voice_styles_dir": str(self.voice_styles_dir), "default_voice": self.default_voice},
            )

        return voice_map, voices

    # ---------------------------------------------------------------------------------
    # Capabilities
    # ---------------------------------------------------------------------------------
    async def get_capabilities(self) -> TTSCapabilities:
        """Return declared Supertonic capabilities."""
        voices = self._voice_infos or []
        return TTSCapabilities(
            provider_name="Supertonic",
            supported_languages=self.SUPPORTED_LANGUAGES,
            supported_voices=voices,
            supported_formats=self.SUPPORTED_FORMATS,
            max_text_length=self.MAX_TEXT_LENGTH,
            supports_streaming=True,
            supports_voice_cloning=False,
            supports_emotion_control=False,
            supports_speech_rate=True,
            supports_pitch_control=False,
            supports_volume_control=False,
            supports_ssml=False,
            supports_phonemes=False,
            supports_multi_speaker=False,
            supports_background_audio=False,
            latency_ms=3500,
            sample_rate=self.sample_rate,
            default_format=AudioFormat.WAV,
        )

    # ---------------------------------------------------------------------------------
    # Generation
    # ---------------------------------------------------------------------------------
    async def generate(self, request: TTSRequest) -> TTSResponse:
        """Generate audio using Supertonic (non-streaming and pseudo-streaming)."""
        if not await self.ensure_initialized():
            raise TTSProviderNotConfiguredError(
                "Supertonic adapter not initialized",
                provider=self.PROVIDER_KEY,
            )

        if request.format not in self.SUPPORTED_FORMATS:
            raise TTSUnsupportedFormatError(
                f"Format {request.format.value} not supported by Supertonic",
                provider=self.PROVIDER_KEY,
            )

        try:
            validate_tts_request(request, provider=self.PROVIDER_KEY)
        except TTSValidationError:
            raise
        except Exception as exc:
            raise TTSValidationError(
                f"Validation failed for Supertonic request: {exc}",
                provider=self.PROVIDER_KEY,
            ) from exc

        voice_id = request.voice or self.default_voice
        style_path = self._voice_to_path.get(voice_id)
        if not style_path:
            raise TTSVoiceNotFoundError(
                f"Voice '{voice_id}' not found for Supertonic",
                provider=self.PROVIDER_KEY,
                details={"voice": voice_id},
            )

        extras = request.extra_params or {}
        total_step = int(extras.get("total_step", self.default_total_step))
        speed = float(request.speed if request.speed is not None else self.default_speed)

        # Load voice style
        try:
            style = await asyncio.to_thread(self._load_voice_style, [str(style_path)], False)
        except Exception as exc:
            raise TTSModelLoadError(
                f"Failed to load Supertonic voice style {style_path}",
                provider=self.PROVIDER_KEY,
                details={"error": str(exc)},
            ) from exc

        # Run inference under lock to protect the engine
        try:
            async with self._engine_lock:
                wav, duration = await asyncio.to_thread(
                    self._engine, request.text, style, total_step, speed
                )
        except Exception as exc:
            raise TTSGenerationError(
                "Supertonic generation failed",
                provider=self.PROVIDER_KEY,
                details={"error": str(exc)},
            ) from exc

        audio_array = self._prepare_audio_array(wav, duration)
        if audio_array.dtype != np.int16:
            audio_array = self._audio_normalizer.normalize(audio_array, target_dtype=np.int16)

        audio_bytes = await self.convert_audio_format(
            audio_array,
            source_format=AudioFormat.PCM,
            target_format=request.format,
            sample_rate=self.sample_rate,
        )

        if request.stream:
            audio_stream = self._build_stream(audio_bytes)
            return TTSResponse(
                audio_stream=audio_stream,
                format=request.format,
                sample_rate=self.sample_rate,
                channels=1,
                text_processed=request.text,
                voice_used=voice_id,
                provider=self.PROVIDER_KEY,
                model=request.model or "tts-supertonic-1",
            )

        return TTSResponse(
            audio_data=audio_bytes,
            format=request.format,
            sample_rate=self.sample_rate,
            channels=1,
            text_processed=request.text,
            voice_used=voice_id,
            provider=self.PROVIDER_KEY,
            model=request.model or "tts-supertonic-1",
        )

    def _prepare_audio_array(self, wav: Any, duration: Any) -> np.ndarray:
        """Trim and normalize engine output to a 1-D numpy array."""
        try:
            arr = np.asarray(wav)
        except Exception as exc:
            raise TTSGenerationError(
                "Supertonic returned non-array audio data",
                provider=self.PROVIDER_KEY,
                details={"error": str(exc)},
            ) from exc

        if arr.ndim == 2:
            if arr.shape[0] != 1:
                raise TTSGenerationError(
                    f"Expected batch size 1 from Supertonic, got shape {arr.shape}",
                    provider=self.PROVIDER_KEY,
                )
            arr = arr[0]
        elif arr.ndim != 1:
            raise TTSGenerationError(
                f"Unexpected Supertonic audio shape {arr.shape}",
                provider=self.PROVIDER_KEY,
            )

        try:
            if isinstance(duration, (list, tuple, np.ndarray)):
                dur_val = float(duration[0])
            else:
                dur_val = float(duration)
            end_idx = int(self.sample_rate * dur_val)
            if end_idx > 0:
                arr = arr[:end_idx]
        except Exception:
            # If duration is unavailable, fall back to full array
            pass
        return arr

    def _build_stream(self, audio_bytes: bytes) -> AsyncGenerator[bytes, None]:
        """Create pseudo-streaming generator from encoded audio bytes."""
        chunk_size = max(1024, int(self.stream_chunk_size or 8192))

        async def _byte_stream():
            for i in range(0, len(audio_bytes), chunk_size):
                chunk = audio_bytes[i:i + chunk_size]
                if chunk:
                    yield chunk

        return _byte_stream()

    async def close(self):
        """Clean up resources and reset state."""
        try:
            await super().close()
        finally:
            self._engine = None
            self._load_voice_style = None
            self._voice_infos = []
            self._voice_to_path = {}
